#!/usr/bin/env python
'''

Created on 04/dic/2012

Project: UNIMORE IMAGE SEARCH
A simple image classification tool

@author:Arletti Luca, Bicocchi Nicola, Spina Andrea

'''

import os
import sys
import math
import array
import time
from datetime import datetime
import numpy
import struct
import pprint

from PIL import Image
from PIL.ExifTags import TAGS

from config import Config
from wordnet_processor import Wordnet_Voter

class Image_Searcher:
    
    def __init__(self):
        numpy.seterr(all='ignore')
        
        self.nn = []
        self.count = 0
        self.start = datetime.now()
        
        self.images = open(Config['IMAGES'], 'rb')
        self.metadata = open(Config['METADATA'],'rb')
        
        self.images.seek(0, 2)
        self.n_images = self.images.tell() / Config['IMGSIZE']
        self.images.seek(0)
        
        self.wordnet = Wordnet_Voter()
        return
    
    def __get_metadata(self, index):
        size = Config['METADATASIZE']
        
        self.metadata.seek(size * index)
        keyword = struct.unpack_from('>80s', self.metadata.read(size))
        return keyword[0].strip()
    
    def __get_exif(self, img):
        info = img._getexif()
        ret = {}
        for tag, value in info.items():
            decoded = TAGS.get(tag, tag)
            ret[decoded] = value
        return ret

    def __thumbnail(self, img):
        """
        Squares the image, resizes it to 32x32, and saves it as filename_n
        """
        
        xsize, ysize = img.size
        if xsize % 2 != 0: xsize -= 1
        if ysize % 2 != 0: ysize -= 1
        box = (0, 0, xsize, ysize)
        
        if xsize < ysize:
            delta = (ysize - xsize) / 2
            box = (0, delta, xsize, ysize - delta)
        else:
            delta = (xsize - ysize) / 2
            box = (delta, 0, xsize - delta, ysize)
        
        img = img.crop(box).resize((32, 32))
        return img
    
    def __normalize(self, img):
        # initializes color channels
        red, green, blue = img[0], img[1], img[2]
        
        # subtracts the mean and divides by stddev each channel
        red_mean, red_std = numpy.mean(red), numpy.std(red)
        green_mean, green_std = numpy.mean(green), numpy.std(green)
        blue_mean, blue_std = numpy.mean(blue), numpy.std(blue)
        
        red = (red - red_mean) / red_std
        green = (green - green_mean) / green_std
        blue = (blue - blue_mean) / blue_std
        
        # makes the array PIL-compatible
        return numpy.dstack((red, green, blue)).reshape((32, 32, 3))
    
    def __image_to_array(self, img):    
        red = array.array('B')
        green = array.array('B')
        blu = array.array('B')
        
        pv = list(img.getdata())
        for pixel in pv:
            red.append(pixel[0])
            green.append(pixel[1])
            blu.append(pixel[2])
            
        red.extend(green)
        red.extend(blu)
        return numpy.array(red, dtype=numpy.uint8).reshape((3, 32, 32))
    
    def __print_stats(self): 
        now = datetime.now()
        ds = (now - self.start).microseconds / float(1000000) + (now - self.start).seconds + (now - self.start).days
   
        print 'started: %s; processed: %d/%d images in %s; total time: %.1f hours' % (
            str(self.start),
            self.count, 
            self.n_images,
            str(now - self.start),
            self.n_images * ds / float(self.count * 3600))
        return
            
    def __wallpaper(self, image, filename):
        raw_imgs = [image['target']]     
        for dist, idx, img, img_norm, keyword in image['dists']: 
            if img is not None:
                raw_imgs.append(img)    
        
        width = 32
        n = int(math.sqrt(len(raw_imgs)) + 1)
        wallpaper = Image.new("RGB", (width * n, width * n))
        for i in xrange(n):
            for j in xrange(n):
                try:
                    raw_img = raw_imgs.pop()
                except IndexError:
                    break
                red, green, blue = raw_img[0], raw_img[1], raw_img[2]
                raw_img = numpy.dstack((red, green, blue)).reshape((32, 32, 3)) 
                raw_img = numpy.rot90(raw_img, 3)                   
                img = Image.fromarray(raw_img)
                wallpaper.paste(img, (width * j, width * i))
            else:
                continue
            break
        wallpaper.save(filename, 'JPEG')            
        return
                
    def load_images(self):
        for filename in os.listdir(Config['TESTSETDIR']): 
            print 'Loading ' + filename + '...'
            try:
                img = Image.open(os.path.join(Config['TESTSETDIR'], filename))
            except IOError:
                continue
            
            try:
                exif_data = self.__get_exif(img)
                timestamp = datetime.fromtimestamp(time.mktime(time.strptime(exif_data['DateTimeOriginal'], "%Y:%m:%d %H:%M:%S")))
            except AttributeError:
                timestamp = None
            
            img = img.transpose(Image.ROTATE_90)
            img = self.__thumbnail(img)
            img = self.__image_to_array(img)

            self.nn.append({
                'filename' : os.path.join(Config['TESTSETDIR'], filename),  
                'timestamp' : timestamp,     
                'keyword' : '',       
                'target' : img,
                'target_norm' : self.__normalize(img),
                'dists' : [(sys.maxint, None, None, None, None)],
                })  
        return len(self.nn)
    
    def search(self, debug=False):
        while True:        
            # read an image from dataset
            try:
                lastimg = numpy.fromfile(self.images, dtype=numpy.uint8, count=Config['IMGSIZE']).reshape((3, 32, 32))
                lastimg_norm = self.__normalize(lastimg)
            except ValueError:
                break
            
            for image in self.nn:
                # extracts an item from testset images
                img_norm = image['target_norm'] 
                dists = image['dists']    
                            
                # analyzing an image
                d = numpy.linalg.norm(img_norm - lastimg_norm)
                if d < dists[len(dists) - 1][0]:
                    dists.append((d, self.count, lastimg, lastimg_norm, self.__get_metadata(self.count)))
                    dists.sort()
                    
                # keeping the list in shape
                if len(dists) > Config['NNN']:
                    dists.pop()
                    
            # partial results
            if self.count % 1000000 == 0 and self.count != 0:
                self.__print_stats()
                if debug:     
                    for r in self.get_results():
                        print '%20s,%30s,%30s' % (r[0], r[1], r[2])
            self.count += 1
        return
    
    def get_results(self):
        results = []
        for image in self.nn:
            keywords = []
            for neighbor in image['dists']:
                keywords.append(neighbor[4])
            image['keyword'] = self.wordnet.vote(keywords)    
            
            filename = Config['RESULTSDIR'] + os.path.basename(image['filename']) + '_' + str(self.count) + '_graph.pdf'
            self.wordnet.draw_graph(filename)
                
            filename = Config['RESULTSDIR'] + os.path.basename(image['filename']) + '_' + str(self.count) + '_wallpaper.jpg'
            self.__wallpaper(image, filename)    
            
            results.append((image['timestamp'], image['keyword'], image['filename']))
        return results
    
if __name__ == '__main__':
    
    try:
        import psyco
        psyco.full()
    except ImportError:
        pass
    
    ims = Image_Searcher()
    ims.load_images()
    ims.search(debug=True)    
    for r in ims.get_results():
        print '%20s,%30s,%30s' % (r[0], r[1], r[2])


    sys.exit()